import os
import datetime

#打印时间
def printbar():
    nowtime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    print("\n"+"=========="*8 + "%s"%nowtime)

#mac系统上pytorch和matplotlib在jupyter中同时跑需要更改环境变量
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

import torch
import torchvision
from torch import nn
from torch.utils.data import Dataset,DataLoader
import torchvision.transforms as transforms
from torchvision import transforms,datasets

transform_train = transforms.Compose(
    [transforms.ToTensor()])
transform_valid = transforms.Compose(
    [transforms.ToTensor()])

ds_train = datasets.ImageFolder("/home/nvidia/cifar2/train/",
            transform = transform_train,target_transform= lambda t:torch.tensor([t]).float())
ds_valid = datasets.ImageFolder("//home/nvidia/cifar2/test/",
            transform = transform_train,target_transform= lambda t:torch.tensor([t]).float())

print(ds_train.class_to_idx)

BATCH_SIZE = 4
dl_train = DataLoader(ds_train,batch_size = BATCH_SIZE,shuffle = True,num_workers=1)
dl_valid = DataLoader(ds_valid,batch_size = BATCH_SIZE,shuffle = True,num_workers=1)
classes = ('red', 'blue', 'yellow','other')

import numpy as np

#查看部分样本
from matplotlib import pyplot as plt
def imshow(img):
    img = img / 2 + 0.5
    np_img = img.numpy()
    print(np_img.shape)
    plt.imshow(np.transpose(np_img,(1,2,0)))
    plt.show()
dataiter = iter(dl_train)
imgs, labels = dataiter.next()
labels = labels.numpy()
labels = labels.reshape(labels.shape[0]*labels.shape[1], )
labels = torch.LongTensor(labels)
print(imgs.shape)
print(imgs)
# print(labels.shape)

imshow(torchvision.utils.make_grid(imgs))
# print(labels)

# Pytorch的图片默认顺序是 Batch,Channel,Width,Height
for x,y in dl_train:
    print(x.shape,y.shape)
    break

#测试AdaptiveMaxPool2d的效果
pool = nn.AdaptiveMaxPool2d((1,1))
t = torch.randn(10,8,32,32)
pool(t).shape

import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, 5)
        self.dropout = nn.Dropout2d(p = 0.5)
        self.adaptive_pool = nn.AdaptiveMaxPool2d((1,1))

        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(64, 32)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(32, 4)

    def forward(self,x):
            x = self.conv1(x)
            x = self.pool(x)
            x = self.conv2(x)
            x = self.pool(x)
            x = self.dropout(x)
            x = self.adaptive_pool(x)
            x = self.flatten(x)
            x = self.fc1(x)
            x = self.relu(x)
            x = self.fc2(x)
            return x
net = Net()
print(net)

import torchkeras
torchkeras.summary(net,input_shape= (3,32,32))

import pandas as pd
import torch.optim as optim
from sklearn.metrics import roc_auc_score
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

def get_num_correct(predictions, labels):  #求准确率
    return predictions.eq(labels).sum().item()


import time

for epoch in range(100):  # loop over the dataset multiple times
    t0 = time.clock()
    t_0 = time.time()
    running_loss = 0.0
    running_correct = 0.0
    for i, data in enumerate(dl_train, 0):
        # get the inputs; data is a list of [inputs, labels]
        t1 = time.clock()
        t_1 = time.time()
        #         print('t1-t0',t1-t0)
        net.train()
        inputs, labels = data
        #         print(inputs.shape)
        # zero the parameter gradients
        optimizer.zero_grad()

        labels = labels.numpy()
        labels = labels.reshape(labels.shape[0] * labels.shape[1], )
        labels = torch.LongTensor(labels)

        # forward + backward + optimize
        outputs = net(inputs)
        #         print(outputs)
        loss = criterion(outputs, labels)
        #         print(loss)
        #         print('labels',labels)
        loss.backward()
        optimizer.step()
        t2 = time.clock()
        t_2 = time.time()
        # print("t2-t1", t2 - t1)
        _, predictions = torch.max(outputs, 1)
        #         print(predictions)
        # print statistics
        running_loss += loss.item()
        #         print(running_loss)
        running_correct += get_num_correct(predictions, labels)
        #         print(loss.item())
        #         t2 = time.clock()
        #         t_2 = time.time()
        #         print("t2-t1",t2-t1)
        if i % 100 == 99:  # print every 50 mini-batches
            print('[%d, %5d] loss: %.3f correct: %.3f' %
                  (epoch + 1, i + 1, running_loss / 100, running_correct / 400))
            t3 = time.clock()
            t_3 = time.time()
            #             print("t3-t0",t3-t0)
            nowtime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            print("\n" + "==========" * 8 + "%s" % nowtime)
            running_loss = 0.0

print('Finished Training')

PATH = './training_model'
torch.save(net.state_dict(), PATH)

dataiter = iter(dl_valid)
images, labels = dataiter.next()
labels = labels.numpy()
labels = labels.reshape(labels.shape[0]*labels.shape[1], )
labels = torch.LongTensor(labels)
print(labels)
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))

net = Net()
net.load_state_dict(torch.load(PATH))
outputs = net(images)

_, predicted = torch.max(outputs, 1)
# print(predicted)
# print(outputs)
# print(torch.max(outputs, 1))
print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]
                              for j in range(4)))

classes = ('red', 'blue', 'yellow', 'other')
correct = 0
total = 0
# since we're not training, we don't need to calculate the gradients for our outputs
with torch.no_grad():
    for data in dl_valid:
        images, labels = data
        labels = labels.numpy()
        labels = labels.reshape(labels.shape[0] * labels.shape[1], )
        labels = torch.LongTensor(labels)
        # calculate outputs by running images through the network
        outputs = net(images)
        #         print(outputs.data.shape)
        # the class with the highest energy is what we choose as prediction
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        # print(labels)
        correct += get_num_correct(predicted, labels)

print('Accuracy of the network on the 140 test images: %d %%' % (100 *
                                                                 correct / total))

# prepare to count predictions for each class
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}

# again no gradients needed
with torch.no_grad():
    for data in dl_valid:
        images, labels = data
        labels = labels.numpy()
        labels = labels.reshape(labels.shape[0]*labels.shape[1], )
        labels = torch.LongTensor(labels)
        outputs = net(images)
        _, predictions = torch.max(outputs, 1)
        # collect the correct predictions for each class
        for label, prediction in zip(labels, predictions):
#             label = torch.LongTensor(label)
            if label == prediction:
                correct_pred[classes[label]] += 1
            total_pred[classes[label]] += 1


# print accuracy for each class
for classname, correct_count in correct_pred.items():
    accuracy = 100 * float(correct_count) / total_pred[classname]
    print("Accuracy for class {:5s} is: {:.1f} %".format(classname,
                                                   accuracy))